package com.graphcoloring;

import org.sat4j.core.VecInt;
import org.sat4j.minisat.SolverFactory;
import org.sat4j.specs.ISolver;

import java.util.Arrays;
import java.util.List;

public class ChenAMOSolver implements Solver {
	public Solution solve(GraphInputData data) throws Exception {
		int n = data.getVertexNumber();
		List<Integer[]> edges = data.getEdgeList();

		int colorCount = n;
		int colors[] = new int[n];
		for (int k = n; k >= 1; k--) {
			int normalVariableNumber = n * k + 1;
			
			int p = (int) Math.sqrt(k);
			if (p*p!=k) p++;
			
			int q = k / p + (k % p != 0 ? 1 : 0);
			int auxiliaryVariableNumber = n*(p+q);
			
//			System.out.println("k=" + k + " p=" + p + " q=" + q);
			
			ISolver solver = SolverFactory.newDefault();
			solver.setTimeout(60);
			solver.newVar(normalVariableNumber + auxiliaryVariableNumber);
			
			// normalVariableNumber+1...normalVariableNumber+p - Ui
			// normalVariableNumber+1+p...normalVariableNumber+1+p+q - Vj
			
			// AMO(U)
			for (int i = 0; i < p; i++) {
				for (int j = i+1; j < p; j++) {
					int clause[] = { -(normalVariableNumber+1+i), -(normalVariableNumber+1+j) };
					solver.addClause(new VecInt(clause));
//					System.out.println(Arrays.toString(clause));
				}
			}
			
			// AMO(V)
			for (int i = 0; i < q; i++) {
				for (int j = i+1; j < q; j++) {
					int clause[] = { -(normalVariableNumber+1+i+p), -(normalVariableNumber+1+j+p) };
					solver.addClause(new VecInt(clause));
//					System.out.println(Arrays.toString(clause));
				}
			}

			// This big clause 
			for (int l = 0; l < n; l++) {
				for (int i = 0; i < p; i++) {
					for (int j = 0; j < q; j++) {
						int xk = l*n+i*q+j+1;
//						System.out.println(xk);
						int uClause[] = {-xk, normalVariableNumber+xk+1+i};
						solver.addClause(new VecInt(uClause));
						
						int vClause[] = {-xk, normalVariableNumber+xk+1+j+p};
						solver.addClause(new VecInt(vClause));
						
	//					System.out.println(Arrays.toString(uClause));
	//					System.out.println(Arrays.toString(vClause));
					}
				}
			}

			for (Integer[] edge : edges) {
				int vertex1 = edge[0] + 1;
				int vertex2 = edge[1] + 1;

				for (int j = 0; j < k; j++) {
					int clause[] = { -(vertex1 * k - j), -(vertex2 * k - j) };
					solver.addClause(new VecInt(clause));
				}
			}
			for (int i = 1; i <= n; i++) {
				int clause[] = new int[k];
				for (int j = 0; j < k; j++) {
					clause[j] = i * k - j;
				}
				solver.addClause(new VecInt(clause));
//				System.out.println(Arrays.toString(clause));
			}
			if (solver.isSatisfiable()) {
				colorCount = k;

				int[] model = solver.model();
				for (int i = 0; i < normalVariableNumber; i++) {
					int var = model[i];
					if (var > 0) {
						int vertex = (var - 1) / k;
						int color = (var - 1 + k) % k;
						colors[vertex] = color;
					}
				}
			} else {
				Solution solution = new Solution();
				solution.colorCount = colorCount;
				solution.colors = colors;

				return solution;

			}
		}
		return null;
	}
}
